package com.hadwinling.alogriithm.projectforpso.ppppppp

import org.apache.spark.rdd.RDD
import org.apache.spark.{SparkConf, SparkContext}

import scala.collection.mutable.ArrayBuffer

/*
Step1:初始化一群粒子（粒子个数为20个），包括随即位置和速度；

Step2:计算每个粒子的适应度fitness；

Step3:对每个粒子，将其适应度与其进过的最好位置（局部）pbest做比较，如果较好，则将其作为当前的最好位置pbest；

Step4:对每个粒子，将其将其适应值与群体所经过的最好位置gbest做比较，如果较好，则将其作为当前最好位置gbest；

Step5:更新所有粒子位置和速度；

Step6:未达到结束条件则转Step2。
*/


object PSS extends Serializable {
  var dimension = 2 //设置维度
  var number_of_particle = 1000 //粒子数
  var no_of_iteration_External = 10 //gbest必须跨节点同步的次数
  var no_of_iteration_Internal = 1000 // 每一个粒子都要计算pbest的次数 迭代次数
  var c1 = 2
  var c2 = 2
  var w = 0.5
  var global_best_position = Array.fill(dimension)(math.random) //全局最优位置

  def main(args: Array[String]): Unit = {
    var conf = new SparkConf().setMaster("local").setAppName("mypso")
    var sc = new SparkContext(config = conf)

    //初始化所有粒子
    var swarm = ArrayBuffer[Particle]()
    for (i <- 0 to number_of_particle - 1) {
      swarm += new Particle(dimension) with Serializable //初始化粒子，并加入到数组中
    }

    var swarm_rdd = sc.parallelize(swarm, 8) //分区设置，一般设置为：工作节点数*cpu核心数

    swarm_rdd.foreach(f => init_particle_best(f))  //初始化最佳粒子
    swarm_rdd.foreach(f => f.compute_particle_best_fitness()) //计算最佳适应值
    swarm_rdd.foreach(f => f.compute_particle_fitness()) //计算当前适应值
    swarm_rdd.foreach(f => global_best_position(f)) //计算全局最佳位置


    var temp_rdd: RDD[Particle] = sc.emptyRDD[Particle]
    temp_rdd = swarm_rdd

    var update_velocity_swarm: RDD[Particle] = sc.emptyRDD[Particle]

    var bCash = sc.broadcast(global_best_position) // 将全局最优解设置为广播

    println(fitness(global_best_position)) //得到最初的全局的最优

    for (iteration <- 0 to no_of_iteration_External - 1) {
      update_velocity_swarm = temp_rdd.map { x => update_particle(x) } // 更新每一个粒子
      temp_rdd = update_velocity_swarm
      temp_rdd.collect().foreach(f => global_best_position(f)) //计算新的全局最佳
      println(iteration+":"+fitness(global_best_position)) // 打印出新计算出的全局最优适应值
      bCash.unpersist(blocking = true) //从辅助节点中删除旧的全局最佳
      bCash = sc.broadcast(global_best_position) //广播更新全局最佳
    }

    /**
     * 更新粒子操作
     * @param particle
     * @return
     */
    def update_particle(particle: Particle): Particle = {
      global_best_position = bCash.value //将全局最佳初始化为广播值
      for (iteration <- 0 to no_of_iteration_Internal - 1) { //更新粒子

        for (i <- 0 to dimension - 1) { //更新速度
          var toward_particle_best = c1 * math.random() * (particle.particle_best_position(i) - particle.particle_position(i))
          var toward_global_best = c2 * math.random() * (global_best_position(i) - particle.particle_position(i))
          particle.particle_velocity(i) = w * particle.particle_velocity(i) + toward_global_best + toward_particle_best
        }

        for (i <- 0 to dimension - 1) { //更新位置
          particle.particle_position(i) = particle.particle_position(i) + particle.particle_velocity(i)
        }

        particle.compute_particle_fitness() //评估每个粒子的函数适应值

        if (particle.particle_best_fitness < fitness(global_best_position)) { //更新每个粒子历史最优位置
          global_best_position = particle.particle_position
        }

      }
      return particle
    }


    sc.stop()
  }

  /**
   * 根据当前位置计算适应值
   *
   * @return
   */
  def fitness(x: Array[Double]): Double = {
    var temp: Double = 0
    for (dim <- 0 to x.length - 1) {
      temp = temp + (math.pow(x(dim), 2))
    }
    return temp
  }


  /**
   * 计算全局最优解(根据自身的历史最优位置和全局最优位置，更新每个粒子的速度和位置)
   * @param particle
   */
  def global_best_position(particle: Particle): Particle = {
    if (particle.particle_best_fitness < fitness(global_best_position)) {
      global_best_position = particle.particle_best_position
    }
    return particle
  }

  /**
   * 初始化粒子的最优位置
   * @param particle
   */
  def init_particle_best(particle: Particle) = {
    particle.particle_best_position = particle.particle_position
  }

}
